#include "shm_bw.h"
#include <gtest/gtest.h>
#include <limits>
#include <spdlog/spdlog.h>

bool TestShmBW(size_t N) {
  spdlog::info("Create tensors...");
  float *d_x = nullptr;
  float *d_out = nullptr;

  cudaMalloc(&d_x, N * sizeof(*d_x));
  cudaMalloc(&d_out, N * sizeof(*d_x));

  std::vector<float> x(N, 1.0f);
  // Copy vectors to device
  cudaMemcpy(d_x, x.data(), N * sizeof(*d_x), cudaMemcpyHostToDevice);

  // Perform the addition.
  cudaStream_t stream;
  cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
  shm_bw_ori(d_out, d_x, 1.0, stream, N);
  shm_bw_ori(d_out, d_x, 1.0, stream, N);
  cudaStreamSynchronize(stream);
  std::vector<float> out(N, -1.0f);
  cudaMemcpy(out.data(), d_out, N * sizeof(*d_out), cudaMemcpyDeviceToHost);

  cudaFree(d_x);
  cudaFree(d_out);

  // Verify that returned results are OK.
  float max_err = -std::numeric_limits<float>::infinity();
  for (size_t i = 0; i < N; ++i) {
    max_err = std::max(max_err, std::abs(out[i] - (x[i] + 100.0f)));
  }
  spdlog::info("[shm_bw] Max error: {}", max_err);

  return true;
}

TEST(test_shm_bw, EXAMPLE) { EXPECT_TRUE(TestShmBW(128 * 64 * 8192)); }
